# Copyright (c) Huawei Technologies Co., Ltd. 2023. All rights reserved.
from typing import Union, List
import random
from atk.case_generator.generator.generate_types import GENERATOR_REGISTRY
from atk.case_generator.generator.base_generator import CaseGenerator
from atk.configs.case_config import InputCaseConfig, CaseConfig

@GENERATOR_REGISTRY.register("ascend_generate_aclnn_moe_token_unpermute_probs_none")
class DefaultGenerator(CaseGenerator):
    def __init__(self, config):
        super().__init__(config)
        self.tokens_num = 1
        self.top_k = 1
        self.hidden_size = 1
        self.case_dtype = None

    def after_input_config(
            self,
            index: int,
            input_case: Union[InputCaseConfig, List[InputCaseConfig]],
            ) -> Union[InputCaseConfig, List[InputCaseConfig]]:
        if index == 2:
            self.tokens_num = input_case.shape[0]
            input_case.dtype = self.case_dtype

        if index == 0:
            self.hidden_size = input_case.shape[1]
            self.case_dtype = input_case.dtype
        return input_case
    
    def after_case_config(self, case_config: CaseConfig) -> CaseConfig:
        case_config.inputs[0].shape = [self.tokens_num * self.top_k, self.hidden_size]
        case_config.inputs[1].shape = [self.tokens_num * self.top_k]
        case_config.inputs[2].shape = [self.tokens_num, self.top_k]
        return case_config

